from typing import Optional

import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import base64
import io
from PIL import Image
from paddleocr import PaddleOCR

app = FastAPI()


# 定义请求体模型
class Item(BaseModel):
    image: Optional[str] = None  # 可选的 Base64 编码字符串
    url: Optional[str] = None  # 可选的图片URL地址


# 初始化 OCR 模型
ocr = PaddleOCR(use_angle_cls=True, lang="ch")  # 如果你想识别中文，设置 lang="ch"


@app.post("/recognize_base64/")
async def recognize_base64(item: Item):
    try:
        # 将 Base64 字符串解码为图像
        image_bytes = base64.b64decode(item.image)
        image = Image.open(io.BytesIO(image_bytes))

        # 使用 PaddleOCR 进行识别
        result = ocr.ocr(image, cls=True)

        # 提取识别结果
        text_result = [line[1][0] for line in result[0]]
        return {"recognized_text": text_result}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/recognize_url/")
async def recognize_url(item: Item):
    try:
        # 下载图片
        response = requests.get(item.url)
        response.raise_for_status()  # 确保请求成功

        # 将下载的数据转换为图像
        image = Image.open(io.BytesIO(response.content))

        # 使用 PaddleOCR 进行识别
        result = ocr.ocr(image, cls=True)

        # 提取识别结果
        text_result = [line[1][0] for line in result[0]]
        return {"recognized_text": text_result}

    except requests.RequestException as e:
        raise HTTPException(status_code=400, detail=f"Failed to download image: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
